import gym
from gym import spaces
import numpy as np
from envs.env_core import EnvCore


class ContinuousActionEnv(object):
    """
    对于连续动作环境的封装
    """

    def __init__(self):
        """
        初始化连续动作环境
        """
        self.env = EnvCore()  # 创建核心环境
        self.num_agent = self.env.agent_num  # 智能体数量

        self.signal_obs_dim = self.env.obs_dim  # 单个智能体的观察维度
        self.signal_action_dim = self.env.action_dim  # 单个智能体的动作维度

        # 如果为真，动作是一个0...N的数字，否则动作是一个one-hot的N维向量
        self.discrete_action_input = False

        self.movable = True  # 是否可移动

        # 配置空间
        self.action_space = []  # 动作空间列表
        self.observation_space = []  # 观察空间列表
        self.share_observation_space = []  # 共享观察空间列表

        share_obs_dim = 0  # 共享观察维度
        total_action_space = []  # 总动作空间
        for agent in range(self.num_agent):
            # 物理动作空间
            u_action_space = spaces.Box(
                low=-np.inf,
                high=+np.inf,
                shape=(self.signal_action_dim,),
                dtype=np.float32,
            )  # 连续动作空间，范围为[-inf, inf]

            if self.movable:
                total_action_space.append(u_action_space)  # 添加动作空间

            # 总动作空间
            self.action_space.append(total_action_space[0])  # 添加动作空间

            # 观察空间
            share_obs_dim += self.signal_obs_dim  # 累加共享观察维度
            self.observation_space.append(
                spaces.Box(
                    low=-np.inf,
                    high=+np.inf,
                    shape=(self.signal_obs_dim,),
                    dtype=np.float32,
                )
            )  # [-inf,inf]

        # 共享观察空间
        self.share_observation_space = [
            spaces.Box(
                low=-np.inf, high=+np.inf, shape=(share_obs_dim,), dtype=np.float32
            )
            for _ in range(self.num_agent)
        ]

    def step(self, actions):
        """
        执行环境步进
        
        输入actions维度假设：
        # actions shape = (5, 2, 5)
        # 5个线程的环境，里面有2个智能体，每个智能体的动作是一个one_hot的5维编码
        
        参数:
            actions: 智能体的动作
            
        返回:
            观察、奖励、完成标志和信息
        """
        results = self.env.step(actions)  # 执行环境步进
        obs, rews, dones, infos = results  # 解包结果
        return np.stack(obs), np.stack(rews), np.stack(dones), infos  # 返回结果

    def reset(self):
        """
        重置环境
        
        返回:
            初始观察
        """
        obs = self.env.reset()  # 重置环境
        return np.stack(obs)  # 返回堆叠的观察

    def close(self):
        """
        关闭环境
        """
        pass

    def render(self, mode="rgb_array"):
        """
        渲染环境
        
        参数:
            mode: 渲染模式
        """
        pass

    def seed(self, seed):
        """
        设置随机种子
        
        参数:
            seed: 随机种子
        """
        pass
